# 导入transformers库中的关键组件
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch

"""
我们在自然语言处理（NLP）的开发中，会使用到HuggingFace的transformers库。
"""

# 这里注意，运行时如果本地没有这个模型，程序会自动从HuggingFace的服务器下载，需要保持网络畅通
# 默认会下载到 C:\Users\用户名\.cache\huggingface\ 路径下

# 如果不想保存在C盘，可以考虑使用软连接指向其他盘符。注意执行此命令之前，目标文件夹（后者）可以存在也可以不存在，来源文件夹（前者）必须不存在。(myfox是我的用户名，你要改成你自己的)
# 例如我这里的：mklink /J "C:\Users\myfox\.cache\huggingface\" "E:\ProgramFiles\AIapplication\huggingface_cache\"


class ModelLoader:
    @staticmethod
    def create_pipeline(
        task_type="text-generation", model_name="gpt2", truncation=True
    ):
        """创建HuggingFace Pipeline
        Args:
            task_type: 任务类型，默认为"text-generation"(文本生成)
            model_name: 预训练模型名称，默认为"gpt2"。
                      GPT-2(Generative Pre-trained Transformer 2)是由OpenAI开发的自回归语言模型，
                      专门用于生成类似人类的文本。它是GPT系列的第二代模型，拥有更强的文本生成能力。
            truncation: 是否启用文本截断，默认为True。
                      当输入文本超过模型的最大长度限制时，将自动截断文本。
                      设置为True可以避免相关的警告信息。
        Returns:
            返回一个可以直接使用的NLP任务pipeline
        Note:
            pipeline()是transformers库中最简单的接口，它将模型加载、预处理和后处理打包成一个易用的函数，
            使得即使不了解底层细节，也能快速应用预训练模型完成NLP任务。
        """
        # pipeline(): 创建一个端到端的NLP任务处理流程，自动处理模型加载、输入处理和输出格式化
        return pipeline(task_type, model=model_name, truncation=truncation)

    @staticmethod
    def load_text_generation_model(model_name="gpt2"):
        """加载文本生成模型和分词器，展示手动处理方式
        Args:
            model_name: 预训练模型名称，默认为"gpt2"。
                      GPT-2是一个强大的语言模型，特别适合文本生成任务。
        Returns:
            返回加载的模型和对应的分词器，如果加载失败则返回None, None
        Note:
            这个方法展示了如何手动使用model和tokenizer进行文本生成，
            相比pipeline方式，这种方法提供了更多的控制和自定义选项。
        """
        try:
            # 加载因果语言模型（Causal LM）和分词器
            model = AutoModelForCausalLM.from_pretrained(model_name)
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            return model, tokenizer
        except Exception as e:
            print(f"模型加载失败: {str(e)}")
            return None, None

    @staticmethod
    def generate_text_manual(model, tokenizer, prompt, max_length=200, temperature=0.7):
        """使用model和tokenizer手动生成文本
        Args:
            model: 预训练语言模型
            tokenizer: 分词器
            prompt: 输入提示文本
            max_length: 生成文本的最大长度
            temperature: 生成文本的随机性，值越大生成的文本越随机
        Returns:
            生成的文本
        """
        try:
            # 将输入文本转换为模型可以理解的格式
            inputs = tokenizer(prompt, return_tensors="pt")

            # 设置生成参数
            gen_kwargs = {
                "max_length": max_length,
                "temperature": temperature,
                "do_sample": True,  # 启用采样模式，使temperature参数生效
                "num_return_sequences": 1,
                "pad_token_id": tokenizer.eos_token_id,
            }

            # 生成文本
            with torch.no_grad():  # 不计算梯度，节省内存
                output_sequences = model.generate(**inputs, **gen_kwargs)

            # 解码生成的文本
            generated_text = tokenizer.decode(
                output_sequences[0], skip_special_tokens=True
            )
            return generated_text
        except Exception as e:
            print(f"文本生成失败: {str(e)}")
            return None


"""
model/tokenizer与pipeline的区别
1. 复杂度：
   - pipeline：高度封装，一行代码完成从文本输入到结果输出的全过程
   - model+tokenizer：需要手动处理tokenization、模型输入、输出解析等步骤
2. 灵活性：
   - pipeline：适合快速应用，但自定义选项有限
   - model+tokenizer：可以完全控制处理流程的每一步，适合研究和自定义任务
3. 性能：
   - pipeline：可能包含额外的处理步骤，在某些场景下性能略低
   - model+tokenizer：可以优化每一步，潜在性能更好
4. 适用场景：
   - pipeline：适合快速原型开发、简单应用场景
   - model+tokenizer：适合研究、复杂任务、需要自定义处理逻辑的场景
在实际应用中，如果需要快速实现功能，pipeline是最佳选择；如果需要更多控制或自定义处理逻辑，则使用model和tokenizer的方式更合适。
"""

if __name__ == "__main__":
    # 示例用法
    # 1. 使用pipeline方式生成文本
    print("=== 使用Pipeline方式生成文本 ===")
    generator = ModelLoader.create_pipeline()
    pipeline_result = generator("Artificial Intelligence will", max_length=200)
    print("\nPipeline生成结果:", pipeline_result[0]["generated_text"])

    # 2. 使用model+tokenizer方式生成文本
    print("\n=== 使用Model+Tokenizer方式生成文本 ===")
    model, tokenizer = ModelLoader.load_text_generation_model()
    if model and tokenizer:
        manual_result = ModelLoader.generate_text_manual(
            model,
            tokenizer,
            "Artificial Intelligence will",
            max_length=200,
            temperature=0.7,
        )
        print("\n手动生成结果:", manual_result)
